#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <bpf/bpf_helpers.h> // for bpf_printk, etc.
#include <bpf/bpf_endian.h> // for bpf_ntohs, etc.

// 定义 PONG 响应的 Redis 协议格式
// +PONG\r\n
#define PONG_RESPONSE "+PONG\r\n"
#define PONG_LEN (sizeof(PONG_RESPONSE) - 1) // -1 for null terminator

// 定义 PING 命令的 Redis 协议格式
// *1\r\n$4\r\nPING\r\n
const char ping_command_prefix[] = "*1\r\n$4\r\nPING\r\n";
#define PING_PREFIX_LEN (sizeof(ping_command_prefix) - 1)

// Redis 默认端口
#define REDIS_PORT 6379

// 辅助函数：计算 TCP/IP 校验和（简化版，仅用于演示）
// 实际生产环境需要更健壮的校验和计算
static __always_inline void calculate_ip_checksum(struct iphdr *iph) {
    iph->check = 0;
    unsigned long sum = 0;
    __u16 *ptr = (__u16 *)iph;
    for (int i = 0; i < sizeof(struct iphdr) / 2; i++) {
        sum += *ptr++;
    }
    sum = (sum & 0xffff) + (sum >> 16);
    sum = (sum & 0xffff) + (sum >> 16);
    iph->check = ~sum;
}

static __always_inline void calculate_tcp_checksum(struct tcphdr *tcph, struct iphdr *iph, void *data, __u16 data_len) {
    tcph->check = 0;
    unsigned long sum = 0;

    // Pseudo-header
    sum += bpf_ntohs(iph->saddr >> 16);
    sum += bpf_ntohs(iph->saddr & 0xFFFF);
    sum += bpf_ntohs(iph->daddr >> 16);
    sum += bpf_ntohs(iph->daddr & 0xFFFF);
    sum += bpf_ntohs(IPPROTO_TCP);
    sum += bpf_ntohs(tcph->len); // TCP header + payload length

    // TCP header
    __u16 *ptr = (__u16 *)tcph;
    for (int i = 0; i < sizeof(struct tcphdr) / 2; i++) {
        sum += *ptr++;
    }

    // TCP payload
    ptr = (__u16 *)data;
    for (int i = 0; i < data_len / 2; i++) {
        sum += *ptr++;
    }
    if (data_len % 2 == 1) { // Odd length payload
        sum += ((__u8 *)data)[data_len - 1];
    }

    sum = (sum & 0xffff) + (sum >> 16);
    sum = (sum & 0xffff) + (sum >> 16);
    tcph->check = ~sum;
}


SEC("xdp_redis_offload")
int xdp_redis_ping_offload(struct xdp_md *ctx) {
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;

    // 1. 以太网头
    struct ethhdr *eth = data;
    if (eth + 1 > data_end)
        return XDP_PASS; // 不是有效的以太网帧

    // 2. IP 头
    struct iphdr *iph = data + sizeof(struct ethhdr);
    if (iph + 1 > data_end)
        return XDP_PASS; // 不是有效的 IP 包

    if (iph->protocol != IPPROTO_TCP)
        return XDP_PASS; // 不是 TCP 包

    // 3. TCP 头
    struct tcphdr *tcph = data + sizeof(struct ethhdr) + sizeof(struct iphdr);
    if (tcph + 1 > data_end)
        return XDP_PASS; // 不是有效的 TCP 段

    // 检查目标端口是否是 Redis 端口
    if (bpf_ntohs(tcph->dest) != REDIS_PORT) {
        return XDP_PASS; // 不是发往 Redis 端口的包
    }

    // 4. TCP Payload (Redis 命令)
    void *payload = data + sizeof(struct ethhdr) + sizeof(struct iphdr) + (tcph->doff * 4);
    __u16 payload_len = bpf_ntohs(iph->tot_len) - sizeof(struct iphdr) - (tcph->doff * 4);

    if (payload + PING_PREFIX_LEN > data_end || payload_len < PING_PREFIX_LEN) {
        return XDP_PASS; // 负载太短，无法是 PING 命令
    }

    // 检查是否是 PING 命令 (简单的前缀匹配)
    // 实际需要更健壮的 Redis 协议解析，但此处为简化演示
    if (bpf_memcmp(payload, ping_command_prefix, PING_PREFIX_LEN) != 0) {
        return XDP_PASS; // 不是 PING 命令，放行给 Redis 处理
    }

    // 5. 调整数据包大小以适应 PONG 响应
    // 原数据包大小: 负载长度
    // 新数据包大小: PONG_LEN
    // Delta: PONG_LEN - payload_len
    int new_pkt_len = sizeof(struct ethhdr) + sizeof(struct iphdr) + (tcph->doff * 4) + PONG_LEN;
    if (bpf_xdp_adjust_head(ctx, -new_pkt_len)) { // 调整包头，使数据区能容纳新包
        return XDP_ABORTED; // 调整失败
    }
    // 重新获取数据指针，因为 bpf_xdp_adjust_head 可能会移动包头
    data = (void *)(long)ctx->data;
    data_end = (void *)(long)ctx->data_end;
    eth = data;
    iph = data + sizeof(struct ethhdr);
    tcph = data + sizeof(struct ethhdr) + sizeof(struct iphdr);
    payload = data + sizeof(struct ethhdr) + sizeof(struct iphdr) + (tcph->doff * 4);


    // 6. 填充 PONG 响应
    if (payload + PONG_LEN > data_end) {
        return XDP_ABORTED; // 空间不足
    }
    bpf_memcpy(payload, PONG_RESPONSE, PONG_LEN);

    // 7. 修改以太网头 (交换源/目的 MAC 地址)
    unsigned char tmp_mac[ETH_ALEN];
    bpf_memcpy(tmp_mac, eth->h_source, ETH_ALEN);
    bpf_memcpy(eth->h_source, eth->h_dest, ETH_ALEN);
    bpf_memcpy(eth->h_dest, tmp_mac, ETH_ALEN);

    // 8. 修改 IP 头 (交换源/目的 IP 地址，更新长度和校验和)
    __be32 tmp_ip = iph->saddr;
    iph->saddr = iph->daddr;
    iph->daddr = tmp_ip;
    iph->tot_len = bpf_htons(sizeof(struct iphdr) + (tcph->doff * 4) + PONG_LEN);
    calculate_ip_checksum(iph);

    // 9. 修改 TCP 头 (交换源/目的端口，更新长度和校验和)
    __be16 tmp_port = tcph->source;
    tcph->source = tcph->dest;
    tcph->dest = tmp_port;
    tcph->len = bpf_htons(sizeof(struct tcphdr) + PONG_LEN); // TCP header len + payload len
    tcph->seq = bpf_ntohl(tcph->seq) + payload_len; // 确认已消费的字节数
    tcph->ack_seq = bpf_ntohl(tcph->ack_seq); // 通常不需要修改，除非是 TCP ACK
    calculate_tcp_checksum(tcph, iph, payload, PONG_LEN);

    // 10. 返回 XDP_TX，指示内核直接发送此修改后的数据包
    return XDP_TX;
}

char _license[] SEC("license") = "GPL";